import threading
import msvcrt
import serial
import struct
import time
import sys
import os
from tqdm import tqdm

# 串口参数定义
PORT = "COM3"
BUADRATE = 4000000
TIMEOUT = 2
# 命令字节定义
ERASE_CHIP = 0xABCDEF01
PROGRAM = 0xABCDEF02
READ_CHIP = 0xABCDEF03
PROGRAM_EXIT = 0xABCDEF04
WRITE_ZK_OK_FLAG = 0xABCDEF05
GET_FLASH_ID = 0xABCDEF06
# 握手命令
START_RECEIVE_TASK = 0x00
TASK_FINISH = 0x01
READ_PRAMA = 0x02
ERROR = 0x03
START_SEND_DATA = 0x04
TASK_BUSY = 0x05
# 文件定义
origin_bin = "GB2312_80_new.bin"
check_bin = "test.bin"
ser = serial.Serial(PORT, BUADRATE, timeout=TIMEOUT)


def count_byte_differences():
    # 打开两个文件以进行二进制读取
    with open(origin_bin, "rb") as file1, open(check_bin, "rb") as file2:
        content1 = file1.read()
        content2 = file2.read()
        # 确保两个文件具有相同的长度
        if len(content1) != len(content2):
            print("文件长度不同，无法比较字节差异, 写入不成功")
            return
        # 逐个字节比较文件内容，并计算差异的字节数
        byte_diff_count = sum(b1 != b2 for b1, b2 in zip(content1, content2))
    if byte_diff_count == 0:
        print(f"比较了读取的数据,数据完全相同, 校验成功!!!!")
    else:
        print(f"Flash中数据与原文件有 {byte_diff_count} 字节的差异, 校验不成功!!!!")


# STM32会发送START_RECEIVE_TASK后,等待100ms接收命令
def wait_start_sign():
    ser.timeout = 0.01
    while ser.read(10):
        pass
    while True:
        data = ser.read(1)
        if len(data) == 1 and data[0] == START_RECEIVE_TASK:
            return


# STM32需要数据时会发出相应的信号,根据信号来传输数据
def wait_sign(sign: int, timeout: float = 0.1, msg: str = "Timeout Error"):
    ser.timeout = timeout
    data = ser.read(1)
    if len(data) == 1 and data[0] == sign:
        return
    else:
        tqdm.write(msg + f" wait sign = {sign} timeout = {timeout}")
        sys.exit()


# 长时间等待信号,用于清空芯片时提示
def wait_sign_with_busy_ack(sign: int, timeout: float = 0.2, msg: str = "等待忙信号标志时发生错误"):
    ser.timeout = timeout
    while True:
        data = ser.read(1)
        if len(data) == 1 and data[0] == sign:
            return
        elif len(data) == 1 and data[0] == TASK_BUSY:
            print(".", end="")
        else:
            print("\n" + msg + f" wait sign = {sign} timeout = {timeout}")
            sys.exit()


def write_zk_OK_flag():
    print("准备向字库写入字库已写好标志")
    wait_start_sign()
    ser.write(struct.pack("<I", WRITE_ZK_OK_FLAG))
    wait_sign(TASK_FINISH, 0.5, "等待STM32返回写好标志信号时超时了,标志好像没写好!")

def read_size_of_data(addr, size):
    # 同步主机发出信号
    wait_start_sign()
    ser.write(struct.pack("<I", READ_CHIP))
    # 发送地址
    wait_sign(READ_PRAMA, timeout=0.2, msg="等待传送读取数量时超时")
    ser.write(struct.pack("<I", addr))
    # 发送数量
    wait_sign(READ_PRAMA, timeout=0.2, msg="等待传送读取地址时超时")
    ser.write(struct.pack("<I", size))
    # 开始接收
    wait_sign(START_SEND_DATA, timeout=0.2, msg="等待开始发送数据时超时")
    ser.timeout = 1
    data = ser.read(size)
    # 未收到,打印超时
    if(len(data) < size):
        tqdm.write(f"从 {addr} 读取 {size} 个字节时时出现了超时, 只读到了 {len(data)} 个字节数据.")
    return data
    
def read_data(size):
    t_start = time.time()
    bytes_read = 0
    addr = 0
    n = 0
    pbar = tqdm(total=size, ncols=100, unit="B", unit_scale=True)
    tqdm.write(f"正在读取 {size} 字节的芯片内容,保存到 {check_bin} 中")
    with open(check_bin, "wb") as f:
        while True:
            # 是否读取完成
            if size - bytes_read == 0:
                # 打印读取完成
                tqdm.write(f"读取用时: {time.time()-t_start:0.3f},一共读取了 {bytes_read} 个字节")
                return
            # 读4096字节,或者更少
            if size - bytes_read >= 4096:
                n = 4096
            else:
                n = size - bytes_read;
            data = read_size_of_data(addr, n)
            # 自增地址
            addr += n
            bytes_read += n
            # 写入文件
            f.write(data)
            # 更新进度条
            pbar.update(n)
            


def erase_chip():
    t_start = time.time()
    print(f"正在擦除芯片内容,等待开始")
    wait_start_sign()
    print(f"擦除开始")
    ser.write(struct.pack("<I", ERASE_CHIP))
    wait_sign_with_busy_ack(TASK_FINISH)
    print(f"整片擦除用时: {time.time()-t_start:0.3f}")
    time.sleep(0.2)


def program():
    addr = 0
    t_start = time.time()
    pbar = tqdm(total=os.path.getsize(origin_bin), ncols=100, unit="B", unit_scale=True)
    tqdm.write("正在将字库写入Flash")
    with open(origin_bin, "rb") as file:
        while True:
            # 发送下载命令
            wait_start_sign()
            ser.write(struct.pack("<I", PROGRAM))
            # 发送地址
            wait_sign(READ_PRAMA)
            ser.write(struct.pack("<I", addr))
            # 读取文件
            data = file.read(4096)
            # 地址递增,并检查是否需要继续写入
            if len(data) != 0:
                addr += len(data)
                pbar.update(len(data))
            else:
                break
            # 发送size
            wait_sign(READ_PRAMA)
            ser.write(struct.pack("<I", len(data)))
            # 发送数据
            wait_sign(READ_PRAMA)
            ser.write(data)
            # 等待完成
            wait_sign(TASK_FINISH, 0.2)


def program_exit():
    # 发送命令
    ser.write(struct.pack("<I", PROGRAM_EXIT))
    wait_sign("Timeout")

def read_flash_id():
    print("准备读取FLASH ID,等待主机发出开始传输消息信号")
    wait_start_sign();
    ser.write(struct.pack("<I", GET_FLASH_ID))
    wait_sign(START_SEND_DATA, msg="等待发送FLASH ID时超时了")
    ret = ser.read(2)
    if len(ret) == 2:
        print(f"你所使用的FLASH型号为 {ret} ")

def count_sequence(data: bytes, sequence: bytes) -> int:  
    count = 0  
    length = len(sequence)  
    for i in range(len(data) - length + 1):  
        if data[i:i+length] == sequence:  
            count += 1  
    return count 

if __name__ == "__main__":
    os.system("cls")
    read_flash_id();
    erase_chip()
    program()
    read_data(os.path.getsize(origin_bin))
    write_zk_OK_flag()
    count_byte_differences()


